"""技术指标计算助手模块"""

import pandas as pd
import numpy as np
from typing import List, Dict, Optional
import logging
from src.utils.logger_manager import log, err

logger = logging.getLogger(__name__)

class TechnicalIndicators:
    """技术指标计算助手类"""

    @staticmethod
    def calculate_ma(df: pd.DataFrame, periods: List[int] = [5, 10, 20, 30, 60]) -> pd.DataFrame:
        """计算移动平均线 (MA)
        
        滚动周期: 每个周期对应 periods 参数中的值 (如 5 日、10 日等)
        
        Args:
            df: 包含OHLCV数据的DataFrame,需包含'close'和'date'列
            periods: MA周期列表,默认[5,10,20,30,60]天
            
        Returns:
            添加了MA列的DataFrame,列名格式为'MA5','MA10'等
        """
        try:
            # 检查必要的列
            if 'close' not in df.columns:
                err("数据中缺少close列,无法计算MA")
                return df
                
            result = df.copy()
            
            # 确保按日期排序
            if 'date' in result.columns:
                result = result.sort_values('date')
            
            # 填充空值
            result['close'] = result['close'].ffill()
            
            # 计算各周期MA
            for period in periods:
                col_name = f'MA{period}'  # MA5, MA10等
                result[col_name] = result['close'].rolling(
                    window=period,
                    min_periods=1  # 允许计算第一个窗口期的值
                ).mean()
                
                # 处理首行的NaN值
                result[col_name] = result[col_name].fillna(result['close'])
                
            return result
            
        except Exception as e:
            logger.error(f"计算MA失败: {str(e)}")
            return df

    @staticmethod
    def calculate_macd(df: pd.DataFrame, fast_period: int = 12, slow_period: int = 26, signal_period: int = 9) -> pd.DataFrame:
        """计算MACD指标
        
        滚动周期:
        - 快线 (EMA_fast): 12 日
        - 慢线 (EMA_slow): 26 日
        - 信号线 (SIGNAL): 9 日
        
        Args:
            df: 包含收盘价数据的DataFrame,需包含'close'和'date'列
            fast_period: 快线周期,默认12
            slow_period: 慢线周期,默认26
            signal_period: 信号线周期,默认9
        """
        try:
            if 'close' not in df.columns:
                err("数据中缺少close列,无法计算MACD")
                return df
                
            result = df.copy()
            # 确保按日期排序
            if 'date' in result.columns:
                result = result.sort_values('date')
            
            # 填充空值
            result['close'] = result['close'].ffill()
            
            # 计算快线EMA和慢线EMA
            result['EMA_fast'] = result['close'].ewm(span=fast_period, adjust=False).mean()
            result['EMA_slow'] = result['close'].ewm(span=slow_period, adjust=False).mean()
            
            # 计算MACD线  〖公式〗DIF = EMA(close,12) - EMA(close,26)
            result['MACD'] = result['EMA_fast'] - result['EMA_slow']
            
            # 计算信号线 〖公式〗DEM = EMA(DIF,9)
            result['SIGNAL'] = result['MACD'].ewm(span=signal_period, adjust=False).mean()
            
            # 计算柱状图 〖公式〗OSC = DIF - DEM
            result['HIST'] = result['MACD'] - result['SIGNAL']
            
            # 确保结果精度与平台一致
            result['MACD'] = result['MACD'].round(6)
            result['SIGNAL'] = result['SIGNAL'].round(6)
            result['HIST'] = result['HIST'].round(6)
            
            return result
            
        except Exception as e:
            err(f"计算MACD失败: {str(e)}")
            return df

    @staticmethod
    def calculate_kdj(df: pd.DataFrame, n: int = 9, m1: int = 3, m2: int = 3) -> pd.DataFrame:
        """计算KDJ指标
        
        滚动周期:
        - RSV: 9 日
        - K 值: 3 日平滑
        - D 值: 3 日平滑
        
        Args:
            df: 包含OHLC数据的DataFrame,需包含'high','low','close'和'date'列
            n: RSV周期,默认9
            m1: K值周期,默认3
            m2: D值周期,默认3
        """
        try:
            for col in ['high', 'low', 'close']:
                if col not in df.columns:
                    err(f"数据中缺少{col}列,无法计算KDJ")
                    return df
                    
            result = df.copy()
            # 确保按日期排序
            if 'date' in result.columns:
                result = result.sort_values('date')
            
            # 填充空值
            result['high'] = result['high'].ffill()
            result['low'] = result['low'].ffill()
            result['close'] = result['close'].ffill()
            
            # 确保 high, low, close 列为 float 类型
            result['high'] = result['high'].astype(float)
            result['low'] = result['low'].astype(float)
            result['close'] = result['close'].astype(float)
            
            # 计算RSV
            low_min = result['low'].rolling(window=n, min_periods=1).min()
            high_max = result['high'].rolling(window=n, min_periods=1).max()
            
            # 避免除零错误
            denom = high_max - low_min
            denom = denom.replace(0, np.nan)
            
            rsv = 100 * ((result['close'] - low_min) / denom)
            rsv = rsv.fillna(50)  # 处理除零结果
            
            # 计算KDJ
            result['K'] = rsv.ewm(alpha=1/m1, min_periods=1, adjust=False).mean()
            result['D'] = result['K'].ewm(alpha=1/m2, min_periods=1, adjust=False).mean()
            result['J'] = 3 * result['K'] - 2 * result['D']
            
            return result
            
        except Exception as e:
            err(f"计算KDJ失败: {str(e)}")
            return df

    @staticmethod
    def calculate_rsi(df: pd.DataFrame, periods: List[int] = [6, 12, 24]) -> pd.DataFrame:
        """计算RSI指标

        滚动周期: 每个周期对应 periods 参数中的值 (如 6 日、12 日等)

        Args:
            df: 包含收盘价数据的DataFrame,需包含'close'列
            periods: RSI周期列表,默认[6,12,24]天
        """
        try:
            # 检查必要的列
            if 'close' not in df.columns:
                err("数据中缺少close列,无法计算RSI")
                return df

            result = df.copy()
            result['close'] = result['close'].ffill()  # 填充空值

            # 计算价格变化
            diff = result['close'].diff()

            for period in periods:
                # 计算上涨和下跌
                gain = diff.where(diff > 0, 0)
                loss = -diff.where(diff < 0, 0)

                # 使用指数移动平均计算平均上涨和下跌
                avg_gain = gain.ewm(alpha=1/period, adjust=False).mean()
                avg_loss = loss.ewm(alpha=1/period, adjust=False).mean()

                # 避免除零错误
                rs = avg_gain / avg_loss.replace(0, np.nan)

                # 计算RSI
                result[f'RSI{period}'] = 100 - (100 / (1 + rs))

                # 填充NaN值
                result[f'RSI{period}'] = result[f'RSI{period}'].fillna(50)

            return result
        except Exception as e:
            err(f"计算RSI失败: {str(e)}")
            return df

    @staticmethod
    def calculate_bollinger_bands(df: pd.DataFrame, period: int = 20, std_dev: float = 2) -> pd.DataFrame:
        """计算布林带指标 (Bollinger Bands)
        
        滚动周期: 20 日 (默认)
        
        Args:
            df: 包含收盘价数据的DataFrame,需包含'close'列
            period: 滚动窗口周期,默认20天
            std_dev: 标准差倍数,默认2倍
        """
        try:
            result = df.copy()
            
            # 计算中轨线（20日移动平均线）
            result['BB_middle'] = df['close'].rolling(window=period).mean()
            
            # 计算标准差
            bb_std = df['close'].rolling(window=period).std()
            
            # 计算上轨和下轨
            result['BB_upper'] = result['BB_middle'] + (bb_std * std_dev)
            result['BB_lower'] = result['BB_middle'] - (bb_std * std_dev)
            
            return result
        except Exception as e:
            err(f"计算布林带失败: {str(e)}")
            return df

    @staticmethod
    def calculate_momentum(df: pd.DataFrame, periods: List[int] = [6, 12]) -> pd.DataFrame:
        """计算动量指标 (Momentum)
        
        滚动周期: 每个周期对应 periods 参数中的值 (如 6 日、12 日)
        
        Args:
            df: 包含收盘价数据的DataFrame,需包含'close'列
            periods: 动量指标周期列表,默认[6,12]天
        """
        try:
            result = df.copy()
            for period in periods:
                col_name = f'Momentum{period}'
                result[col_name] = df['close'] - df['close'].shift(period)
            return result
        except Exception as e:
            err(f"计算Momentum失败: {str(e)}")
            return df

    @staticmethod
    def calculate_atr(df: pd.DataFrame, period: int = 14) -> pd.DataFrame:
        """计算ATR指标 (Average True Range)
        
        滚动周期: 14 日 (默认)
        
        Args:
            df: 包含OHLC数据的DataFrame,需包含'high','low','close'列
            period: 滚动窗口周期,默认14天
        """
        try:
            result = df.copy()
            high_low = df['high'] - df['low']
            high_close = np.abs(df['high'] - df['close'].shift())
            low_close = np.abs(df['low'] - df['close'].shift())
            
            ranges = pd.concat([high_low, high_close, low_close], axis=1)
            true_range = np.max(ranges, axis=1)
            
            result['ATR'] = true_range.rolling(window=period).mean()
            return result
        except Exception as e:
            err(f"计算ATR失败: {str(e)}")
            return df

    @staticmethod
    def calculate_obv(df: pd.DataFrame) -> pd.DataFrame:
        """计算OBV指标 (On Balance Volume)
        
        滚动周期: 累积计算，无固定周期
        
        Args:
            df: 包含收盘价和成交量数据的DataFrame,需包含'close'和'volume'列
        """
        try:
            result = df.copy()
            price_diff = df['close'].diff()
            volume = df['volume']
            
            obv = (price_diff > 0) * volume - (price_diff < 0) * volume
            result['OBV'] = obv.cumsum()
            
            return result
        except Exception as e:
            err(f"计算OBV失败: {str(e)}")
            return df

    @staticmethod
    def calculate_williams_r(df: pd.DataFrame, period: int = 14) -> pd.DataFrame:
        """计算威廉指标 (Williams %R)
        
        滚动周期: 14 日 (默认)
        一般取14天用法：1.低于20，可能超买见顶，可考虑卖出
                       2.高于80，可能超卖见底，可考虑买进
                       3.与RSI、MTM指标配合使用，效果更好
        Args:
            df: 包含OHLC数据的DataFrame,需包含'high','low','close'列
            period: 滚动窗口周期,默认14天
        """
        try:
            for col in ['high', 'low', 'close']:
                if col not in df.columns:
                    err(f"数据中缺少{col}列,无法计算威廉指标")
                    return df
            
            result = df.copy()
            # 确保按日期排序
            if 'date' in result.columns:
                result = result.sort_values('date')
            
            # 填充空值
            result['high'] = result['high'].ffill()
            result['low'] = result['low'].ffill()
            result['close'] = result['close'].ffill()
            
            # 计算最高价和最低价的滚动值
            high_max = result['high'].rolling(window=period, min_periods=1).max()
            low_min = result['low'].rolling(window=period, min_periods=1).min()
            
            # 计算威廉指标
            result['Williams_R'] = -100 * ((high_max - result['close']) / (high_max - low_min))
            
            return result
        except Exception as e:
            err(f"计算威廉指标失败: {str(e)}")
            return df

    @staticmethod
    def calculate_cci(df: pd.DataFrame, period: int = 14) -> pd.DataFrame:
        """计算CCI指标 (Commodity Channel Index)
        
        滚动周期: 14 日 (默认)
        
        Args:
            df: 包含OHLC数据的DataFrame,需包含'high','low','close'列
            period: 滚动窗口周期,默认14天
        
        Returns:
            添加了'CCI'列的DataFrame
        """
        try:
            for col in ['high', 'low', 'close']:
                if col not in df.columns:
                    err(f"数据中缺少{col}列,无法计算CCI")
                    return df

            result = df.copy()
            # 计算典型价格 (TP)
            result['TP'] = (result['high'] + result['low'] + result['close']) / 3

            # 计算移动平均价格 (SMA)
            result['SMA'] = result['TP'].rolling(window=period).mean()

            # 计算均绝对偏差 (MAD) 替换为手动计算
            result['MAD'] = result['TP'].rolling(window=period).apply(
                lambda x: np.mean(np.abs(x - np.mean(x))), raw=True
            )

            # 计算CCI
            result['CCI'] = (result['TP'] - result['SMA']) / (0.015 * result['MAD'])

            return result
        except Exception as e:
            err(f"计算CCI失败: {str(e)}")
            return df

    @staticmethod
    def calculate_sar(df: pd.DataFrame, af_start: float = 0.02, af_step: float = 0.02, af_max: float = 0.2) -> pd.DataFrame:
        """计算SAR指标 (Stop and Reverse)
        
        滚动周期: 动态计算，无固定周期
        
        Args:
            df: 包含OHLC数据的DataFrame,需包含'high','low'列
            af_start: 加速因子初始值，默认0.02
            af_step: 加速因子步长，默认0.02
            af_max: 加速因子最大值，默认0.2
        
        Returns:
            添加了'SAR'列的DataFrame
        """
        try:
            for col in ['high', 'low']:
                if col not in df.columns:
                    err(f"数据中缺少{col}列,无法计算SAR")
                    return df

            result = df.copy()
            result['SAR'] = np.nan

            # 初始化参数
            af = af_start
            ep = result['low'][0]  # 极值点，初始为第一天的最低价
            sar = result['high'][0]  # 初始SAR为第一天的最高价
            uptrend = True  # 初始为上涨趋势

            for i in range(1, len(result)):
                prev_sar = sar
                if uptrend:
                    sar = prev_sar + af * (ep - prev_sar)
                    sar = min(sar, result['low'][i - 1], result['low'][i])
                    if result['high'][i] > ep:
                        ep = result['high'][i]
                        af = min(af + af_step, af_max)
                    if result['low'][i] < sar:
                        uptrend = False
                        sar = ep
                        ep = result['low'][i]
                        af = af_start
                else:
                    sar = prev_sar + af * (ep - prev_sar)
                    sar = max(sar, result['high'][i - 1], result['high'][i])
                    if result['low'][i] < ep:
                        ep = result['low'][i]
                        af = min(af + af_step, af_max)
                    if result['high'][i] > sar:
                        uptrend = True
                        sar = ep
                        ep = result['high'][i]
                        af = af_start

                result.at[i, 'SAR'] = sar

            return result
        except Exception as e:
            err(f"计算SAR失败: {str(e)}")
            return df

    @staticmethod
    def calculate_mfi(df: pd.DataFrame, period: int = 14) -> pd.DataFrame:
        """计算MFI指标 (Money Flow Index)
        
        滚动周期: 14 日 (默认)
        
        Args:
            df: 包含OHLCV数据的DataFrame,需包含'high','low','close','volume'列
            period: 滚动窗口周期,默认14天
        """
        try:
            for col in ['high', 'low', 'close', 'volume']:
                if col not in df.columns:
                    err(f"数据中缺少{col}列,无法计算MFI")
                    return df

            result = df.copy()
            # 填充空值
            result['high'] = result['high'].ffill()
            result['low'] = result['low'].ffill()
            result['close'] = result['close'].ffill()
            result['volume'] = result['volume'].ffill()

            # 计算典型价格 (TP)
            result['TP'] = (result['high'] + result['low'] + result['close']) / 3

            # 计算资金流量 (Raw Money Flow)
            result['Raw_MF'] = result['TP'] * result['volume']

            # 计算正负资金流量
            result['Positive_MF'] = np.where(result['TP'] > result['TP'].shift(1), result['Raw_MF'], 0)
            result['Negative_MF'] = np.where(result['TP'] < result['TP'].shift(1), result['Raw_MF'], 0)

            # 计算资金流量比率 (Money Flow Ratio)
            positive_mf_sum = result['Positive_MF'].rolling(window=period).sum()
            negative_mf_sum = result['Negative_MF'].rolling(window=period).sum()
            result['MFR'] = positive_mf_sum / negative_mf_sum

            # 计算MFI
            result['MFI'] = 100 - (100 / (1 + result['MFR']))

            return result
        except Exception as e:
            err(f"计算MFI失败: {str(e)}")
            return df

    @staticmethod
    def calculate_adl(df: pd.DataFrame) -> pd.DataFrame:
        """计算ADL指标 (Accumulation/Distribution Line)
        
        滚动周期: 累积计算，无固定周期
        
        Args:
            df: 包含OHLCV数据的DataFrame,需包含'high','low','close','volume'列
        """
        try:
            for col in ['high', 'low', 'close', 'volume']:
                if col not in df.columns:
                    err(f"数据中缺少{col}列,无法计算ADL")
                    return df

            result = df.copy()
            # 填充空值
            result['high'] = result['high'].ffill()
            result['low'] = result['low'].ffill()
            result['close'] = result['close'].ffill()
            result['volume'] = result['volume'].ffill()

            # 计算CLV (Close Location Value)
            result['CLV'] = ((result['close'] - result['low']) - (result['high'] - result['close'])) / (result['high'] - result['low'])
            result['CLV'] = result['CLV'].fillna(0)  # 处理分母为0的情况

            # 计算ADL
            result['ADL'] = (result['CLV'] * result['volume']).cumsum()

            return result
        except Exception as e:
            err(f"计算ADL失败: {str(e)}")
            return df

    def analyze_data(self, history_data: pd.DataFrame) -> pd.DataFrame:
        """分析历史数据并计算所有技术指标
        
        Args:
            history_data: 包含历史数据的DataFrame,必须包含OHLCV数据和date列
            
        Returns:
            添加了所有技术指标的DataFrame
        """
        try:
            # 检查必要的列
            required_columns = ['date', 'open', 'high', 'low', 'close', 'volume']
            
            if 'date' not in history_data.columns:
                history_data['date'] = history_data['end_date']
                
            missing_columns = [col for col in required_columns if col not in history_data.columns]
            if missing_columns:
                err(f"数据缺少必要的列: {missing_columns}")
                return history_data
                
            # 确保数据按日期排序
            df = history_data.copy()
            df = df.sort_values('date')
            
            # 填充空值
            for col in required_columns[1:]: 
                df[col] = df[col].fillna(method='ffill')

            # 计算技术指标
            df = self.calculate_ma(df)
            df = self.calculate_macd(df)
            df = self.calculate_kdj(df)
            df = self.calculate_rsi(df)
            df = self.calculate_bollinger_bands(df)
            df = self.calculate_momentum(df)
            df = self.calculate_atr(df)
            df = self.calculate_obv(df)
            df = self.calculate_williams_r(df)
            df = self.calculate_cci(df)
            df = self.calculate_sar(df)            
            df = self.calculate_mfi(df)
            df = self.calculate_adl(df)
            
            return df
            
        except Exception as e:
            err(f"计算技术指标失败: {str(e)}")
            return history_data
